package com.apex.spark.transform;

import com.apex.spark.SparkEnvironment;
import com.apex.spark.utils.ConfigKeyName;
import com.typesafe.config.Config;
import org.apache.commons.lang3.StringUtils;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;
import com.apex.spark.transform.SqlCommandParser.SqlCommandCall;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.List;
import java.util.Objects;

/**
 * spark SQL 执行入口
 */
public class Sql implements SparkTransform {
    private final Logger logger = LoggerFactory.getLogger(Sql.class);
    //配置文件管理
    private Config config;
    //用户自定义查询SQL结果集
    private String sql;
    private SparkSession sparkSession;
    //返回结果集
    private Dataset<Row> dataset = null;

    /**
     * 数据处理方法
     * @param env 上下文件执行环境
     * @param data 源端数据
     * @return Dataset<Row>
     */
    @Override
    public Dataset<Row> process(SparkEnvironment env, Dataset<Row> data) {
        List<String> sqlAll = null;
        if (!Objects.isNull(env.getSqlAll())){
            sqlAll = env.getSqlAll();
        }
        sparkSession = env.getSparkSession();
        if (!sqlAll.isEmpty()){
            List<SqlCommandCall> calls = SqlCommandParser.parse(sqlAll);
            for (SqlCommandCall call : calls) {
                logger.info("\n"+call.operands[0]);
                callCommand(call);
            }
        }else if (!StringUtils.isBlank(sql)){
            dataset = sparkSession.sql(sql);
        }else {
            return null;
        }
        return dataset;
    }

    /**
     *
     * @param plugin 执行环境
     */
    @Override
    public void prepare(SparkEnvironment plugin) {
        if (config.hasPath(ConfigKeyName.SPARK_RUN_SQL)) {
            sql = config.getString(ConfigKeyName.SPARK_RUN_SQL);
        } else {
            sql = "";
        }

    }

    /**
     *
     * @return config
     */
    @Override
    public Config getConfig() {
        return config;
    }

    /**
     *
     * @param config config
     */
    @Override
    public void setConfig(Config config) {
        this.config = config;
    }

    /**
     * 分段执行SQL脚本
     * @param cmdCall
     */
    private void callCommand(SqlCommandCall cmdCall) {
        switch (cmdCall.command) {
            case SET:
                callSet(cmdCall);
                break;
            case CREATE_TABLE:
                callCreateTable(cmdCall);
                break;
            case INSERT_INTO:
                callInsertInto(cmdCall);
                break;
            case SELECT_TABLE:
                dataset = callSelect(cmdCall);
                break;
            default:
                throw new RuntimeException("Unsupported command: " + cmdCall.command);
        }
    }

    /**
     * 用户自定义调优命令
     * @param cmdCall
     */
    private void callSet(SqlCommandCall cmdCall) {
        String key = cmdCall.operands[0];
        String value = cmdCall.operands[1];
        logger.info("set " + key + "=" + value + ";");
        sparkSession.conf().set(key, value);
    }

    /**
     * 执行create语句
     * @param cmdCall
     */
    private void callCreateTable(SqlCommandCall cmdCall) {
        String ddl = cmdCall.operands[0];
        logger.info(ddl);
        try {
            sparkSession.sql(ddl);
        } catch (Exception e) {
            throw new RuntimeException("SQL parse failed:\n" + ddl + "\n", e);
        }
    }

    /**
     * 执行insert语句
     * @param cmdCall
     */
    private void callInsertInto(SqlCommandCall cmdCall) {
        String dml = cmdCall.operands[0];
        logger.info(dml);
        try {
            sparkSession.sql(dml);
        } catch (Exception e) {
            throw new RuntimeException("SQL parse failed:\n" + dml + "\n", e);
        }
    }

    /**
     * 执行查询语句并返回查询结果给下游算子
     * @param cmdCall
     * @return Dataset<Row>
     */
    private Dataset<Row> callSelect(SqlCommandCall cmdCall) {
        String dml = cmdCall.operands[0];
        logger.info(dml);
        try {
           return sparkSession.sql(dml);
        } catch (Exception e) {
            throw new RuntimeException("SQL parse failed:\n" + dml + "\n", e);
        }
    }
}
